from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
import os
from langchain.chat_models import init_chat_model

"""
这是classification 的一个小样例
"""
#  从环境变量中读取DeepSeek的API Key
key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

llm = init_chat_model(
    model="gpt-4o-mini",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)

tagging_prompt = ChatPromptTemplate.from_template(
    """
Extract the desired information from the following passage.

Only extract the properties mentioned in the 'Classification' function.

Passage:
{input}
"""
)


# pydantic 中的basemodel
class Classification(BaseModel):
    sentiment: str = Field(description="The sentiment of the text")
    aggressiveness: int = Field(
        description="How aggressive the text is on a scale from 1 to 10"
    )
    language: str = Field(description="The language the text is written in")


# Structured LLM
structured_llm = llm.with_structured_output(schema=Classification)

# 运行结果
# inp = "Estoy increiblemente contento de haberte conocido! Creo que seremos muy buenos amigos!"
# prompt = tagging_prompt.invoke({"input": inp})
# response = structured_llm.invoke(prompt)
#
# print(response)

inp = "Estoy muy enojado con vos! Te voy a dar tu merecido!"
prompt = tagging_prompt.invoke({"input": inp})
response = structured_llm.invoke(prompt)

print(response.model_dump())
